"""
无人机通信可视化模块 - 实现仿真结果的可视化
"""

import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.patches import Circle
import networkx as nx
import pandas as pd
import seaborn as sns
from mpl_toolkits.mplot3d import Axes3D
import matplotlib
import matplotlib.font_manager as fm

# 尝试查找中文字体并加载
font_path = None
chinese_fonts = ['SimHei', 'Microsoft YaHei', 'SimSun', 'STSong', 'KaiTi', 'FangSong', 'Arial Unicode MS']
for font_name in chinese_fonts:
    try:
        font_path = fm.findfont(fm.FontProperties(family=font_name), fallback_to_default=False)
        print(f"找到字体: {font_name}, 路径: {font_path}")
        break
    except:
        continue

if font_path:
    # 设置字体
    plt.rcParams['font.family'] = ['sans-serif']
    plt.rcParams['font.sans-serif'] = [font_name] + plt.rcParams['font.sans-serif']
    plt.rcParams['axes.unicode_minus'] = False
else:
    # 如果找不到中文字体，使用英文标题替代
    print("警告: 未找到中文字体，将使用英文标题")
    plt.rcParams['font.family'] = ['sans-serif']
    plt.rcParams['font.sans-serif'] = ['Arial', 'Helvetica']

class UAVNetworkVisualizer:
    """无人机网络可视化类"""
    
    def __init__(self, sim_results, output_dir='./results'):
        """
        初始化可视化器
        
        参数:
            sim_results: 仿真结果
            output_dir: 输出目录
        """
        self.sim_results = sim_results
        self.output_dir = output_dir
        
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        # 提取数据
        self.events = sim_results['events']
        self.statistics = sim_results['statistics']
        self.uavs = sim_results['uavs']
        self.protocols = sim_results['protocols']
        
        # 修复protocols键名，如果需要
        if 'fly_protocol' in self.protocols:
            self.protocols['dcdp_protocol'] = self.protocols.pop('fly_protocol')
        
        # 修复statistics键名，如果需要
        if 'fly_protocol' in self.statistics:
            self.statistics['dcdp_protocol'] = self.statistics.pop('fly_protocol')
        
        # 传统协议列表
        self.traditional_protocols = ['traditional', 'zigbee', 'wifi_direct', 'lora', 'fanet']
        
        # 每个协议的显示名称和颜色
        self.protocol_display = {
            'traditional': {'name': 'Basic Traditional', 'color': '#3498db'},
            'zigbee': {'name': 'ZigBee', 'color': '#2ecc71'},
            'wifi_direct': {'name': 'WiFi Direct', 'color': '#9b59b6'},
            'lora': {'name': 'LoRa', 'color': '#f1c40f'},
            'fanet': {'name': 'FANET', 'color': '#e67e22'},
            'dcdp_protocol': {'name': 'DCDP Protocol', 'color': '#e74c3c'}
        }
        
        # 配置绘图样式
        sns.set_style("whitegrid")
        
    def plot_protocol_comparison(self):
        """绘制协议性能比较图"""
        # 提取性能数据
        trad = self.statistics['traditional']
        dcdp = self.statistics['dcdp_protocol']
        
        # 计算各项指标
        metrics = {
            'Average Delay (ms)': [trad['avg_delay'], dcdp['avg_delay']],
            'Message Success Rate (%)': [
                trad['successful_messages'] / max(1, trad['total_messages']) * 100,
                dcdp['successful_messages'] / max(1, dcdp['total_messages']) * 100
            ],
            'Emergency Message Rate (%)': [
                trad['emergency_messages_handled'] / max(1, trad['total_messages']) * 100,
                dcdp['emergency_messages_handled'] / max(1, dcdp['total_messages']) * 100
            ],
            'Deadline Met Rate (%)': [
                100 - (trad['missed_deadlines'] / max(1, trad['successful_messages']) * 100),
                100 - (dcdp['missed_deadlines'] / max(1, dcdp['successful_messages']) * 100)
            ]
        }
        
        # 英文标题映射（如果中文无法显示）
        metrics_en = {
            'Average Delay (ms)': 'Average Delay (ms)',
            'Message Success Rate (%)': 'Message Success Rate (%)',
            'Emergency Message Rate (%)': 'Emergency Message Processing Rate (%)',
            'Deadline Met Rate (%)': 'Deadline Met Rate (%)'
        }
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        axes = axes.flatten()
        
        # 绘制每个指标的对比图
        for i, (metric_name, values) in enumerate(metrics.items()):
            ax = axes[i]
            protocols = ['Traditional Protocol', 'DCDP Protocol']
            protocols_en = ['Traditional Protocol', 'DCDP Protocol']
            colors = ['#3498db', '#e74c3c']
            
            # 计算提升百分比
            if i == 0:  # 时延越小越好
                improvement = ((values[0] - values[1]) / values[0]) * 100
                ax.text(1.1, values[1], f"↓ {improvement:.1f}%", fontsize=12, color='green')
            else:  # 其他指标越大越好
                improvement = ((values[1] - values[0]) / max(0.1, values[0])) * 100
                ax.text(1.1, values[1], f"↑ {improvement:.1f}%", fontsize=12, color='green')
            
            # 绘制条形图
            bars = ax.bar(protocols_en, values, color=colors, width=0.6)
            
            # 设置标题和标签
            ax.set_title(metrics_en[metric_name], fontsize=14)
            ax.set_ylim(0, max(values) * 1.2)
            
            # 添加数值标签
            for bar in bars:
                height = bar.get_height()
                ax.text(bar.get_x() + bar.get_width()/2., height + 0.05 * max(values),
                        f'{height:.2f}', ha='center', va='bottom', fontsize=12)
        
        # 整体标题
        plt.suptitle('Performance Comparison: DCDP Protocol vs Traditional Protocol', fontsize=18)
        plt.tight_layout(rect=[0, 0, 1, 0.95])
        
        # 添加注释说明DCDP是什么
        fig.text(0.5, 0.01, 'DCDP: Delay-aware and Dual-Priority UAV Network Protocol', 
                ha='center', fontsize=10, style='italic')
        
        # 保存图表
        plt.savefig(os.path.join(self.output_dir, 'protocol_comparison.png'), dpi=300)
        plt.close()
    
    def plot_delay_distribution(self):
        """绘制时延分布图"""
        # 提取时延数据
        trad_delays = []
        dcdp_delays = []
        
        for event in self.events:
            if event['type'] == 'message_sent' and event['success'] and event['delay'] is not None:
                if event['protocol'] == 'traditional':
                    trad_delays.append(event['delay'])
                else:
                    dcdp_delays.append(event['delay'])
        
        # 创建图表
        plt.figure(figsize=(12, 6))
        
        # 绘制分布
        bins = np.linspace(0, max(max(trad_delays, default=100), max(dcdp_delays, default=100)), 30)
        plt.hist(trad_delays, bins=bins, alpha=0.7, label='Traditional Protocol', color='#3498db')
        plt.hist(dcdp_delays, bins=bins, alpha=0.7, label='DCDP Protocol', color='#e74c3c')
        
        # 添加垂直线表示平均值
        plt.axvline(np.mean(trad_delays), color='#3498db', linestyle='dashed', linewidth=2,
                   label=f'Traditional Protocol Avg: {np.mean(trad_delays):.2f}ms')
        plt.axvline(np.mean(dcdp_delays), color='#e74c3c', linestyle='dashed', linewidth=2,
                   label=f'DCDP Protocol Avg: {np.mean(dcdp_delays):.2f}ms')
        
        # 设置标题和标签
        plt.title('Message Delay Distribution', fontsize=16)
        plt.xlabel('Delay (ms)', fontsize=12)
        plt.ylabel('Message Count', fontsize=12)
        plt.legend(fontsize=12)
        plt.grid(True, alpha=0.3)
        
        # 保存图表
        plt.tight_layout()
        plt.savefig(os.path.join(self.output_dir, 'delay_distribution.png'), dpi=300)
        plt.close()
    
    def plot_success_rate_by_message_type(self):
        """绘制不同消息类型的成功率对比"""
        # 按消息类型统计
        type_stats = {
            'traditional': {msg_type: {'total': 0, 'success': 0} for msg_type in ['EMERGENCY', 'CONTROL', 'TELEMETRY', 'VIDEO', 'DATA']},
            'dcdp_protocol': {msg_type: {'total': 0, 'success': 0} for msg_type in ['EMERGENCY', 'CONTROL', 'TELEMETRY', 'VIDEO', 'DATA']}
        }
        
        # 统计数据
        for event in self.events:
            if event['type'] == 'message_sent' and 'message_type' in event:
                protocol = event['protocol']
                msg_type = event['message_type']
                
                if msg_type in type_stats[protocol]:
                    type_stats[protocol][msg_type]['total'] += 1
                    if event['success']:
                        type_stats[protocol][msg_type]['success'] += 1
        
        # 计算成功率
        success_rates = {
            'traditional': {},
            'dcdp_protocol': {}
        }
        
        for protocol in ['traditional', 'dcdp_protocol']:
            for msg_type, stats in type_stats[protocol].items():
                if stats['total'] > 0:
                    success_rates[protocol][msg_type] = (stats['success'] / stats['total']) * 100
                else:
                    success_rates[protocol][msg_type] = 0
        
        # 创建数据框
        msg_types = list(success_rates['traditional'].keys())
        trad_rates = [success_rates['traditional'][t] for t in msg_types]
        dcdp_rates = [success_rates['dcdp_protocol'][t] for t in msg_types]
        
        # 创建图表
        plt.figure(figsize=(12, 6))
        x = np.arange(len(msg_types))
        width = 0.35
        
        plt.bar(x - width/2, trad_rates, width, label='Traditional Protocol', color='#3498db')
        plt.bar(x + width/2, dcdp_rates, width, label='DCDP Protocol', color='#e74c3c')
        
        # 添加数值标签
        for i, v in enumerate(trad_rates):
            plt.text(i - width/2, v + 1, f'{v:.1f}%', ha='center', fontsize=10)
        
        for i, v in enumerate(dcdp_rates):
            plt.text(i + width/2, v + 1, f'{v:.1f}%', ha='center', fontsize=10)
            # 添加提升标签
            if trad_rates[i] > 0:
                improvement = ((v - trad_rates[i]) / trad_rates[i]) * 100
                plt.text(i, max(v, trad_rates[i]) + 5, f'↑{improvement:.1f}%', ha='center', color='green', fontsize=10)
        
        # 设置图表属性
        plt.xlabel('Message Type', fontsize=12)
        plt.ylabel('Success Rate (%)', fontsize=12)
        plt.title('Success Rate by Message Type', fontsize=16)
        plt.xticks(x, msg_types)
        plt.legend(fontsize=12)
        plt.ylim(0, 105)
        
        # 保存图表
        plt.tight_layout()
        plt.savefig(os.path.join(self.output_dir, 'success_rate_by_type.png'), dpi=300)
        plt.close()
    
    def plot_network_topology_snapshot(self):
        """绘制网络拓扑快照"""
        # 创建图
        G = nx.Graph()
        
        # 添加节点
        for uav in self.uavs:
            G.add_node(uav['id'], pos=(uav['position'][0], uav['position'][1]))
        
        # 添加边 (通信链接)
        for uav in self.uavs:
            for neighbor_id in uav['neighbors']:
                G.add_edge(uav['id'], neighbor_id)
        
        # 获取节点位置
        pos = nx.get_node_attributes(G, 'pos')
        
        # 绘制图形
        plt.figure(figsize=(12, 12))
        
        # 绘制节点
        nx.draw_networkx_nodes(G, pos, node_size=500, node_color='#3498db', alpha=0.8)
        
        # 绘制边
        nx.draw_networkx_edges(G, pos, width=1, alpha=0.5, edge_color='gray')
        
        # 绘制标签
        nx.draw_networkx_labels(G, pos, font_size=12, font_family='sans-serif')
        
        # 添加无人机通信范围圆圈
        ax = plt.gca()
        for uav in self.uavs:
            circle = Circle((uav['position'][0], uav['position'][1]), 
                          radius=200, alpha=0.1, color='#95a5a6')
            ax.add_patch(circle)
        
        # 设置图表属性
        plt.title('UAV Communication Network Topology', fontsize=16)
        plt.axis('off')
        
        # 保存图表
        plt.tight_layout()
        plt.savefig(os.path.join(self.output_dir, 'network_topology.png'), dpi=300)
        plt.close()
    
    def create_performance_summary(self):
        """创建性能总结图"""
        # 提取性能对比数据
        performance = {
            'Average Delay (ms)': {
                'Traditional': self.statistics['traditional']['avg_delay'],
                'DCDP Protocol': self.statistics['dcdp_protocol']['avg_delay'],
                'Improvement': ((self.statistics['traditional']['avg_delay'] - self.statistics['dcdp_protocol']['avg_delay']) / 
                        max(0.1, self.statistics['traditional']['avg_delay'])) * 100
            },
            'Message Success Rate (%)': {
                'Traditional': self.statistics['traditional']['successful_messages'] / max(1, self.statistics['traditional']['total_messages']) * 100,
                'DCDP Protocol': self.statistics['dcdp_protocol']['successful_messages'] / max(1, self.statistics['dcdp_protocol']['total_messages']) * 100,
                'Improvement': ((self.statistics['dcdp_protocol']['successful_messages'] / max(1, self.statistics['dcdp_protocol']['total_messages']) - 
                        self.statistics['traditional']['successful_messages'] / max(1, self.statistics['traditional']['total_messages'])) * 100)
            },
            'Emergency Message Rate (%)': {
                'Traditional': self.statistics['traditional']['emergency_messages_handled'] / max(1, self.statistics['traditional']['total_messages']) * 100,
                'DCDP Protocol': self.statistics['dcdp_protocol']['emergency_messages_handled'] / max(1, self.statistics['dcdp_protocol']['total_messages']) * 100,
                'Improvement': ((self.statistics['dcdp_protocol']['emergency_messages_handled'] / max(1, self.statistics['dcdp_protocol']['total_messages']) - 
                        self.statistics['traditional']['emergency_messages_handled'] / max(1, self.statistics['traditional']['total_messages'])) * 100)
            },
            'Deadline Met Rate (%)': {
                'Traditional': 100 - (self.statistics['traditional']['missed_deadlines'] / max(1, self.statistics['traditional']['successful_messages']) * 100),
                'DCDP Protocol': 100 - (self.statistics['dcdp_protocol']['missed_deadlines'] / max(1, self.statistics['dcdp_protocol']['successful_messages']) * 100),
                'Improvement': ((100 - (self.statistics['dcdp_protocol']['missed_deadlines'] / max(1, self.statistics['dcdp_protocol']['successful_messages']) * 100)) - 
                        (100 - (self.statistics['traditional']['missed_deadlines'] / max(1, self.statistics['traditional']['successful_messages']) * 100)))
            }
        }
        
        # 创建数据框
        df = pd.DataFrame(performance).T
        
        # 创建图表
        fig, ax = plt.subplots(figsize=(10, 6))
        
        # 设置索引名称
        df.index.name = 'Performance Metrics'
        
        # 绘制表格
        cell_text = []
        for row in range(len(df)):
            cell_text.append([f"{df.iloc[row][col]:.2f}" + ("" if col == 'Improvement' else "") + 
                            ("%" if col == 'Improvement' else "") for col in df.columns])
        
        table = ax.table(cellText=cell_text,
                      rowLabels=df.index,
                      colLabels=df.columns,
                      cellLoc='center',
                      loc='center',
                      bbox=[0.2, 0.2, 0.7, 0.6])
        
        # 设置表格样式
        table.auto_set_font_size(False)
        table.set_fontsize(12)
        table.scale(1.2, 1.5)
        
        # 为改进率列设置颜色
        for row in range(len(df)):
            cell = table[row+1, 2]  # 改进率列
            if df.iloc[row]['Improvement'] > 0:
                cell.set_facecolor('#c8e6c9')  # 绿色表示正向改进
            else:
                cell.set_facecolor('#ffcdd2')  # 红色表示负向改进
        
        # 设置图表属性
        ax.set_title('DCDP Protocol Performance Improvement Summary', fontsize=16)
        ax.axis('off')
        
        # 添加注释说明
        fig.text(0.5, 0.01, 'DCDP: Delay-aware and Dual-Priority UAV Network Protocol', 
                ha='center', fontsize=10, style='italic')
        
        # 保存图表
        plt.tight_layout()
        plt.savefig(os.path.join(self.output_dir, 'performance_summary.png'), dpi=300)
        plt.close()
    
    def plot_all_protocols_comparison(self):
        """绘制所有协议的性能比较图"""
        # 准备数据
        metrics = {
            'Success Rate (%)': [],
            'Average Delay (ms)': [],
            'Emergency Handling (%)': [],
            'Deadline Met Rate (%)': []
        }
        
        protocols = []
        colors = []
        
        # 收集所有协议的数据
        for protocol_name in self.traditional_protocols + ['dcdp_protocol']:
            if protocol_name not in self.statistics or self.statistics[protocol_name]['total_messages'] == 0:
                continue
                
            stats = self.statistics[protocol_name]
            
            # 计算指标
            success_rate = stats['successful_messages'] / max(1, stats['total_messages']) * 100
            emergency_rate = stats['emergency_messages_handled'] / max(1, stats['total_messages']) * 100
            deadline_met = 100 - (stats['missed_deadlines'] / max(1, stats['successful_messages']) * 100)
            
            # 添加数据
            metrics['Success Rate (%)'].append(success_rate)
            metrics['Average Delay (ms)'].append(stats['avg_delay'])
            metrics['Emergency Handling (%)'].append(emergency_rate)
            metrics['Deadline Met Rate (%)'].append(deadline_met)
            
            # 记录协议名称和颜色
            protocols.append(self.protocol_display[protocol_name]['name'])
            colors.append(self.protocol_display[protocol_name]['color'])
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        axes = axes.flatten()
        
        # 绘制每个指标的对比图
        for i, (metric_name, values) in enumerate(metrics.items()):
            ax = axes[i]
            
            # 绘制条形图
            bars = ax.bar(protocols, values, color=colors)
            
            # 设置标题和标签
            ax.set_title(metric_name, fontsize=14)
            ax.set_ylim(0, max(values) * 1.2 if max(values) > 0 else 100)
            
            # 添加数值标签
            for bar in bars:
                height = bar.get_height()
                ax.text(bar.get_x() + bar.get_width()/2., height + 0.02 * max(values),
                        f'{height:.2f}', ha='center', va='bottom', fontsize=10)
            
            # 调整x轴标签
            plt.setp(ax.get_xticklabels(), rotation=30, ha='right')
        
        # 整体标题
        plt.suptitle('Performance Comparison Across All Protocols', fontsize=18)
        plt.tight_layout(rect=[0, 0, 1, 0.95])
        
        # 保存图表
        plt.savefig(os.path.join(self.output_dir, 'all_protocols_comparison.png'), dpi=300)
        plt.close()
    
    def plot_protocol_delay_distribution(self):
        """绘制不同协议的时延分布对比图"""
        # 收集每个协议的时延数据
        delay_data = {}
        
        for event in self.events:
            if event['type'] == 'message_sent' and event['success'] and event['delay'] is not None:
                protocol = event['protocol']
                if protocol not in delay_data:
                    delay_data[protocol] = []
                delay_data[protocol].append(event['delay'])
        
        # 如果没有足够的数据，返回
        if len(delay_data) < 2:
            return
        
        # 创建图表
        plt.figure(figsize=(14, 8))
        
        # 计算所有时延的范围
        all_delays = []
        for delays in delay_data.values():
            all_delays.extend(delays)
        
        max_delay = max(all_delays) if all_delays else 100
        bins = np.linspace(0, max_delay, 30)
        
        # 绘制每个协议的时延分布
        for protocol, delays in delay_data.items():
            if protocol in self.protocol_display and len(delays) > 0:
                display_info = self.protocol_display[protocol]
                plt.hist(delays, bins=bins, alpha=0.6, label=display_info['name'], 
                         color=display_info['color'], edgecolor='black', linewidth=0.5)
                
                # 添加平均值线
                avg_delay = np.mean(delays)
                plt.axvline(avg_delay, color=display_info['color'], linestyle='dashed', linewidth=2,
                           label=f"{display_info['name']} Avg: {avg_delay:.2f}ms")
        
        # 设置标题和标签
        plt.title('Delay Distribution Comparison Across Protocols', fontsize=16)
        plt.xlabel('Delay (ms)', fontsize=12)
        plt.ylabel('Message Count', fontsize=12)
        plt.legend(fontsize=10)
        plt.grid(True, alpha=0.3)
        
        # 保存图表
        plt.tight_layout()
        plt.savefig(os.path.join(self.output_dir, 'protocol_delay_distribution.png'), dpi=300)
        plt.close()
    
    def create_all_visualizations(self):
        """创建所有可视化图表"""
        print("Generating visualization results...")
        
        # 绘制协议性能比较图
        self.plot_protocol_comparison()
        print("- Protocol comparison chart generated")
        
        # 绘制所有协议比较图
        self.plot_all_protocols_comparison()
        print("- All protocols comparison chart generated")
        
        # 绘制协议时延分布对比图
        self.plot_protocol_delay_distribution()
        print("- Protocol delay distribution chart generated")
        
        # 绘制时延分布图
        self.plot_delay_distribution()
        print("- Delay distribution chart generated")
        
        # 绘制不同消息类型的成功率对比
        self.plot_success_rate_by_message_type()
        print("- Success rate by message type chart generated")
        
        # 绘制网络拓扑快照
        self.plot_network_topology_snapshot()
        print("- Network topology chart generated")
        
        # 创建性能总结图
        self.create_performance_summary()
        print("- Performance summary chart generated")
        
        print(f"All visualization results have been saved to {self.output_dir} directory")
        
        return self.output_dir 